# C4.5 Decision Tree 采用信息增益作为最优划分属性的决策树模型
from math import log
import operator as ort
import numpy as np


class DecisionTree:
    # 计算信息熵
    def cal_entropy(self, data):
        """
        :param data: 信息集
        :return: entropy: 信息集的熵
        """
        inst_num = len(data)  # 信息集的样本数量
        labelcounts = {}  # 统计类别出现的次数
        for featVect in data:
            # 特征向量的最后一列为类别标签
            current_label = featVect[-1]
            if current_label not in labelcounts.keys():
                # 有新的类别的时候初始化值为 0
                labelcounts[current_label] = 0  # 让 labelcounts 字典中键 current_label 的值等于 0
            # 统计各个类别的出现次数
            labelcounts[current_label] += 1
        entropy = 0.0
        for key in labelcounts:
            # 计算各个类别出现的频率
            prob = float(labelcounts[key]) / inst_num
            # 计算信息熵
            entropy = prob * log(prob, 2)
        return entropy

    # 划分数据集，根据划分属性和属性的取值对原数据集进行划分
    def split_data(self, data, feature_index, feature_value, feature_type):
        """
        :param data: 数据集
        :param feature_index: 划分属性
        :param feature_value: 划分属性的取值
        :param feature_type: 划分属性类型
        :return: splitedData: 划分后的数据集
        """
        splitedData = []  # 用于保存划分后的数据集子集
        if feature_type == "D":  # 对于离散属性
            for featVect in data:  # featVect: 特征函数
                if featVect[feature_index] == feature_value:
                    reducedFev = featVect[:feature_index]
                    reducedFev = list(reducedFev)
                    reducedFev.extend(list(featVect[feature_index + 1:]))  # extend 扩展函数
                    reducedFev = np.array(reducedFev)
                    splitedData.append(reducedFev)  # 把划分属性剔除，剩余数据放入 splitedData 列表
        if feature_type == "L":  # 对于连续属性 <=
            for featVect in data:
                if featVect[feature_index] <= feature_value:  # 特征向量对应划分属性的值小于设定的阈值
                    splitedData.append(featVect)
        if feature_type == "R":  # 对于连续属性 >
            for featVect in data:
                if featVect[feature_index] > feature_value:
                    splitedData.append(featVect)
        return splitedData

    # 选择最优划分属性
    def choose_best_feature_to_split(self, data):
        """
        :param data: 数据集
        :return: best_feature: 选择的最优划分属性
        """
        feature_num = len(data[0] - 1)  # 可以使用的特征的个数
        baseEnt = self.cal_entropy(data)  # 计算划分前数据集的信息熵
        # 初始化最大信息增益、最优划分属性以及最优划分点
        bestInforGain = 0
        best_feature = -1
        # 计算每个属性的信息增益，选择最大的
        for i in range(feature_num):
            # 离散属性
            featureList = [example[i] for example in data]  # 记录每一个实例的每一个属性
            uniqueValues = set(featureList)  # 把属性拆开存入列表中
            newEnt = 0.0
            for value in uniqueValues:  # 根据没个属性进行数据划分并且计算信息增益
                # 每个离散取值划分一个对应的数据子集
                subdata = self.split_data(data, i, value, "D")
                prob = float(len(subdata)) / len(data)
                # 计算每一种划分的信息熵
                newEnt += prob * self.cal_entropy(subdata)
            # 计算每一种划分的信息增益
            inforGain = newEnt - baseEnt
            # 选取信息增益最大的划分作为划分属性
            if inforGain > bestInforGain:
                bestInforGain = inforGain
                best_feature = i
            # 连续属性的情况暂时未考虑
        return best_feature

    # 统计出现次数最多的类别
    def mayorityCount(self, classList):
        """
        :param classList: 数据集样本类别组成的列表
        :return: sortedclasscount: 出现次数最多的类别
        """
        classcount = {}
        for vote in classList:
            if vote not in classcount.keys():
                classcount[vote] = 0
        classcount[vote] += 1
        # iteritems: 返回字典键值和值作为一个对象
        # operator.itemgetter: 取出对象指定位置上的值
        sortedclasscount = sorted(classcount.iteritems(), ort.itemgetter(1), reverse=True)
        return sortedclasscount[0][0]  # 返回最大的类别

    # 创建决策树
    def create_decision_tree(self, data, labels):
        """
        :param data: 数据集
        :param labels: 属性名的列表
        :return: modelTree: 以字典的形式返回决策树模型
        """
        # 获取数据集的样本类别标签列表
        classList = [example[-1] for example in data]
        # 当类别完全相同的时候停止划分
        if classList.count(classList[0]) == len(classList):  # 剩下的类别完全相同
            return classList[0]
        # 遍历完所有特征时，返回出现次数最多的类别
        if len(data[0]) == 1:
            return self.mayorityCount(classList)
        # 获取最好的分类索引和划分点
        bestFeat = self.choose_best_feature_to_split(data)

        bestFeatLabel = labels[bestFeat]  # 获取该类别的名称
        # 生成一个树节点
        modelTree = {bestFeatLabel: {}}
        del labels[bestFeat]
        # 获取离散属性的所有可能取值
        featValues = [example[bestFeat] for example in data]
        uniqueVals = set(featValues)
        # 给每个值生成一个子节点，通过递归构造决策树
        for value in uniqueVals:
            sublabels = labels[:]
            modelTree[bestFeatLabel][value] = self.create_decision_tree(self.split_data(data, bestFeat, value, "D"),
                                                                        sublabels)
        # 连续属性情况暂不考虑
        return modelTree

    # 决策树预测函数
    def predic(self, tree_model, feature_names, test_vect):
        """
        :param tree_model: 以字典形式保存的决策树模型
        :param feature_names: 特征的名字
        :param test_vect: 测试样本的特征向量
        :return: classlabel: 预测类别标签
        """
        firstStr = list(tree_model.keys())[0]  # 查找根节点
        # 对于离散属性
        secondDict = tree_model[firstStr]  # 根节点的子节点
        featIndex = feature_names.index(firstStr)
        key = test_vect[featIndex]
        valueOfFeat = secondDict[key]
        if isinstance(valueOfFeat, dict):  # 非叶子节点
            classLabel = self.predic(valueOfFeat, feature_names, test_vect)
        else:  # 叶子节点
            classLabel = valueOfFeat
        return classLabel
